from annoy import AnnoyIndex
import numpy as np

# 1. 数据生成
np.random.seed(42)
data = np.random.rand(100, 5)    # 生成100个5维向量
query_point = np.random.rand(1, 5)   # 查询点

# 2. 构建Annoy索引
num_trees = 10    # 构建10棵随机投影树
dimension = data.shape[1]    # 向量维度

# 初始化Annoy索引
annoy_index = AnnoyIndex(dimension, metric='euclidean')  # 使用欧氏距离

# 向索引中添加数据
for i in range(len(data)):
    annoy_index.add_item(i, data[i])

# 构建索引
annoy_index.build(num_trees)

# 3. 查询最近邻
top_k = 5  # 查询最近的5个点
results = annoy_index.get_nns_by_vector(
    query_point[0], top_k, include_distances=True)

# 4. 输出结果
print("查询点:", query_point[0])
print("\n最近邻查询结果:")
for idx, (neighbor_idx, distance) in enumerate(zip(results[0], results[1])):
    print(f"结果 {idx + 1}: 索引 {neighbor_idx}, 距离 {distance:.4f}, 向量 {data[neighbor_idx]}")

# 5. 索引分析
print("\n索引分析:")
print(f"索引包含的点数: {annoy_index.get_n_items()}")
print(f"索引是否已构建: {'是' if annoy_index.on_disk_build else '否'}")